using System.Security.Claims;
using AiPlatformAdmin.Api.Services;
using AiPlatformAdmin.Api.Data;
using Microsoft.EntityFrameworkCore;
using Microsoft.Extensions.DependencyInjection;

namespace AiPlatformAdmin.Api.Middleware;

/// <summary>
/// Middleware to resolve tenant from JWT claims
/// </summary>
public class TenantMiddleware
{
    private readonly RequestDelegate _next;

    public TenantMiddleware(RequestDelegate next)
    {
        _next = next;
    }

    public async Task InvokeAsync(HttpContext context, ITenantService tenantService)
    {
        tenantService.SetCurrentTenantId(null);

        var tenantIdClaim = context.User?.FindFirst("TenantId")?.Value;

        if (!string.IsNullOrEmpty(tenantIdClaim) && int.TryParse(tenantIdClaim, out int tenantId))
        {
            tenantService.SetCurrentTenantId(tenantId);
        }
        else if (context.Request.Headers.TryGetValue("X-Tenant-Id", out var headerTenantId) && int.TryParse(headerTenantId, out var parsedTenantId))
        {
            tenantService.SetCurrentTenantId(parsedTenantId);
        }
        else
        {
            var userIdClaim = context.User?.FindFirst(ClaimTypes.NameIdentifier)?.Value;
            if (!string.IsNullOrEmpty(userIdClaim) && int.TryParse(userIdClaim, out var userId))
            {
                var dbContext = context.RequestServices.GetService<AppDbContext>();
                if (dbContext != null)
                {
                    var resolvedTenantId = await dbContext.Users
                        .AsNoTracking()
                        .Where(u => u.Id == userId)
                        .Select(u => (int?)u.TenantId)
                        .FirstOrDefaultAsync();

                    if (resolvedTenantId.HasValue)
                    {
                        tenantService.SetCurrentTenantId(resolvedTenantId);
                    }
                }
            }
        }

        await _next(context);
    }
}

public static class TenantMiddlewareExtensions
{
    public static IApplicationBuilder UseTenantMiddleware(this IApplicationBuilder builder)
    {
        return builder.UseMiddleware<TenantMiddleware>();
    }
}
